Skip to content

refactor(BA-5710): thread main_access_key through UserPermission#11043

Closed
jopemachine wants to merge 5 commits into
refactor/BA-5650-A-main-access-key-helpersfrom
refactor/BA-5650-B-user-permission-owner-id
Closed

refactor(BA-5710): thread main_access_key through UserPermission#11043
jopemachine wants to merge 5 commits into
refactor/BA-5650-A-main-access-key-helpersfrom
refactor/BA-5650-B-user-permission-owner-id

Conversation

@jopemachine
Copy link
Copy Markdown
Member

@jopemachine jopemachine commented Apr 14, 2026

📚 Stack

Merge from top to bottom. Intermediate slices may not build standalone; the final refactor tip is #11051, and #11040 is the schema-drop follow-up on top of it.

Summary

Rename UserPermission.user_uuidowner_id; add main_access_key: str | None field. KernelRow.to_kernel_info populates it from the eagerly-loaded user_row relationship; search_kernels adds selectinload. KernelRow.delegate_ownership drops access_key arg. recalc_concurrency_used switches to a user_uuid subquery shim.

Resolves BA-5710. Part of epic BA-5650.


📚 Documentation preview 📚: https://sorna--11043.org.readthedocs.build/en/11043/


📚 Documentation preview 📚: https://sorna-ko--11043.org.readthedocs.build/ko/11043/


BA-5650 Series: Split Rationale

Overall goal: migrate the session owner identifier from access_key (keypair) to owner_id (user UUID), and drop the sessions.access_key / kernels.access_key columns.

Split criteria: layer + dependency order. Bottom-up (DB helpers → service → API) so the destructive column drop can land safely at the end.

Order Issue Layer Role
1 BA-5709 Foundation Add get_main_access_key_by_id and related resolver helpers (everything else depends on this)
2 BA-5710 RBAC / Permission Rename UserPermission.user_uuid → owner_id; add main_access_key field
3 BA-5711 Data Rename SessionData.user_uuid → owner_id; Row adapters; GQL node
4 BA-5712 Repository Collapse SessionRepository, SessionDBSource, creators signatures
5 BA-5713 Scheduler predicates / drf / options; scheduler repo; stream/events shims
6 BA-5714 Sokovan allocation / lifecycle / workload, launcher, controller, executor
7 BA-5715 Service Drop owner_id from 21 read/control Actions; resolve via current_user()
8 BA-5716 API (v1) Remove owner_access_key from REST v1 DTOs (breaking)
9 BA-5717 Tests / Cleanup Test suite updates; remaining ORM / gql_legacy touch-ups
10 BA-5653 Schema Alembic migration — rename/drop columns (destructive, must land last)

Why this split

  • Dependencies: the BA-5709 helpers are required before 3–7 can recover access_key from owner_id. BA-5653 (destructive) only runs once every reader has migrated off the dropped column.
  • Review size: a single-shot change would touch hundreds of files. Slicing by layer keeps each PR focused on one concern.
  • Rollback safety: steps 1–7 are no-op on external behavior, step 8 is the API breaking change, step 10 is the schema drop. Each step is independently revertible.

@github-actions github-actions Bot added size:M 30~100 LoC comp:manager Related to Manager component labels Apr 14, 2026
@github-actions github-actions Bot added the area:docs Documentations label Apr 14, 2026
@jopemachine jopemachine changed the title refactor(BA-5650-B): thread main_access_key through UserPermission refactor(BA-5710): thread main_access_key through UserPermission Apr 14, 2026
@jopemachine jopemachine marked this pull request as draft April 14, 2026 07:30
@jopemachine jopemachine force-pushed the refactor/BA-5650-A-main-access-key-helpers branch from 2d7d64e to e66536b Compare April 14, 2026 07:38
@jopemachine jopemachine force-pushed the refactor/BA-5650-B-user-permission-owner-id branch from 6d34e21 to d88a14b Compare April 14, 2026 07:39
@jopemachine jopemachine requested a review from Copilot April 14, 2026 07:46
Comment thread src/ai/backend/manager/models/kernel/row.py
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR continues the BA-5650 refactor stack by changing kernel-level permission data to carry the owning user ID (owner_id) and the user’s main_access_key (nullable), and wiring kernel query paths to load the associated user row so KernelInfo can be populated without relying on kernels.access_key.

Changes:

  • Update UserPermission to use owner_id and main_access_key: str | None.
  • Populate KernelInfo.user_permission.main_access_key from KernelRow.user_row.main_access_key, and eager-load KernelRow.user_row in session kernel search.
  • Update selected unit-test fixtures to the new permission field names, and add a changelog entry.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
tests/unit/manager/sokovan/scheduler/terminator/conftest.py Updates termination fixture data to renamed access-key field(s) and permission fields.
tests/unit/manager/sokovan/scheduler/handlers/conftest.py Updates scheduler handler fixtures for renamed permission/access-key fields and session owner naming.
tests/unit/manager/services/session/test_session_service.py Updates UserPermission construction in a kernel info test helper.
tests/unit/manager/api/compute_sessions/test_handler.py Updates session/kernel test factories to renamed owner/access-key fields.
src/ai/backend/manager/repositories/session/db_source/db_source.py Eager-loads KernelRow.user_row during kernel search to support new KernelInfo mapping.
src/ai/backend/manager/models/kernel/row.py Drops access_key from delegate_ownership, maps KernelInfo permissions via user_row.main_access_key, and changes concurrency recalculation filtering.
src/ai/backend/manager/data/kernel/types.py Renames UserPermission fields to owner_id + main_access_key.
changes/11043.misc.md Adds a changelog note describing the refactor intent and behavior.
Comments suppressed due to low confidence (7)

tests/unit/manager/sokovan/scheduler/handlers/conftest.py:114

  • SessionMetadata (from ai.backend.manager.data.session.types) currently has user_uuid and access_key fields. This fixture now passes owner_id and omits access_key, which will raise a TypeError when constructing SessionMetadata. This PR only changes UserPermission, so the session metadata fields should remain aligned with SessionMetadata until its rename lands.
        metadata=SessionMetadata(
            name=f"session-{sid}",
            domain_name="default",
            group_id=group_id,
            user_uuid=user_uuid,
            access_key=access_key,
            session_type=session_type,
            priority=0,
            created_at=now,
            tag=None,

tests/unit/manager/sokovan/scheduler/handlers/conftest.py:629

  • SessionDataForStart still uses access_key and user_uuid fields (see sokovan/data/lifecycle.py). This factory passes main_access_key and owner_id, so it will raise a TypeError. Keep the argument names aligned with the current dataclass (or move the type rename into this PR).
        sessions_for_start = [
            SessionDataForStart(
                session_id=s.session_info.identity.id,
                creation_id=s.session_info.identity.creation_id,
                access_key=AccessKey(s.session_info.metadata.access_key),
                session_type=s.session_info.identity.session_type,
                name=s.session_info.identity.name,
                cluster_mode=ClusterMode(s.session_info.resource.cluster_mode),
                kernels=[
                    KernelBindingData(
                        kernel_id=KernelId(k.id),
                        agent_id=AgentId(k.resource.agent) if k.resource.agent else None,
                        agent_addr=k.resource.agent_addr,
                        scaling_group=s.session_info.resource.scaling_group_name or "default",
                        image="test-image:latest",
                        architecture=k.image.architecture or "x86_64",
                    )
                    for k in s.kernel_infos
                ],
                user_uuid=s.session_info.metadata.user_uuid,
                user_email="test@example.com",
                user_name="test-user",

tests/unit/manager/sokovan/scheduler/handlers/conftest.py:663

  • TerminatingSessionData still defines the field as access_key, not main_access_key (see repositories/scheduler/types/session.py). This factory will fail when constructing TerminatingSessionData with the renamed argument.
        return [
            TerminatingSessionData(
                session_id=s.session_info.identity.id,
                access_key=AccessKey(s.session_info.metadata.access_key),

tests/unit/manager/api/compute_sessions/test_handler.py:75

  • SessionData (from ai.backend.manager.data.session.types) still uses user_uuid and includes an access_key field. This test factory now passes owner_id and omits access_key, which will raise a TypeError when constructing SessionData. Either keep the factory aligned with the current SessionData signature, or include the SessionData rename in the same PR.
    return SessionData(
        id=session_id or uuid4(),
        session_type=SessionTypes.INTERACTIVE,
        priority=0,
        is_preemptible=True,
        cluster_mode=ClusterMode.SINGLE_NODE,
        cluster_size=1,
        domain_name="default",
        group_id=uuid4(),
        user_uuid=uuid4(),
        occupying_slots=ResourceSlot({"cpu": Decimal("2.0"), "mem": Decimal("4294967296")}),
        requested_slots=ResourceSlot({"cpu": Decimal("4.0"), "mem": Decimal("8589934592")}),

tests/unit/manager/sokovan/scheduler/terminator/conftest.py:142

  • TerminatingSessionData (from ai.backend.manager.repositories.scheduler.types.session) still defines the field as access_key, not main_access_key. This fixture will fail to instantiate the dataclass. Update the fixture to use the current field name (or update the underlying type in the same PR if the rename is intended here).
    return TerminatingSessionData(
        session_id=session_id or SessionId(uuid4()),
        access_key=AccessKey("test-key"),
        creation_id=str(uuid4()),
        status=SessionStatus.TERMINATING,

tests/unit/manager/sokovan/scheduler/handlers/conftest.py:548

  • ScheduledSessionData still defines the field as access_key, not main_access_key (see repositories/scheduler/types/results.py). This factory will fail to construct ScheduledSessionData until the type rename happens.
        scheduled_sessions = [
            ScheduledSessionData(
                session_id=s.session_info.identity.id,
                creation_id=s.session_info.identity.creation_id,
                access_key=AccessKey(s.session_info.metadata.access_key),

tests/unit/manager/sokovan/scheduler/handlers/conftest.py:567

  • SessionDataForPull still defines the field as access_key, not main_access_key (see sokovan/data/lifecycle.py). This factory will fail to instantiate SessionDataForPull with the renamed argument.
        sessions_for_pull = [
            SessionDataForPull(
                session_id=s.session_info.identity.id,
                creation_id=s.session_info.identity.creation_id,
                access_key=AccessKey(s.session_info.metadata.access_key),

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread src/ai/backend/manager/models/kernel/row.py
Comment thread src/ai/backend/manager/models/kernel/row.py
Comment thread src/ai/backend/manager/models/kernel/row.py
Comment thread src/ai/backend/manager/data/kernel/types.py
Rename UserPermission.user_uuid to owner_id and add main_access_key:
str | None. KernelRow.to_kernel_info populates the new field from the
eagerly-loaded user_row relationship; search_kernels now selectinload's
user_row so the scalar is available without N+1.

KernelRow.delegate_ownership drops its access_key argument (resolved
from the owner). recalc_concurrency_used switches from
KernelRow.access_key to a KernelRow.user_uuid subquery keyed by the
owner's main_access_key as a shim for the eventual access_key column
drop.

Test fixtures that construct UserPermission are updated.
Earlier the branch pulled test fixtures that also referenced later
slices' renames (SessionMetadata.owner_id, ScheduledSessionData
.main_access_key, SessionDataForPull/Start.main_access_key,
SessionData.owner_id). Revert those fixtures to main state and
re-apply only the UserPermission field renames (user_uuid -> owner_id,
access_key -> main_access_key) which are in this slice's scope.

Also update the two remaining live readers of UserPermission:
- sokovan/scheduler/fair_share/aggregator.py
- api/adapters/session.py (kernel_info_to_node)

And fix the SessionRow.delegate_ownership caller of
KernelRow.delegate_ownership to match the new single-argument signature.
@jopemachine jopemachine force-pushed the refactor/BA-5650-B-user-permission-owner-id branch from 5ea1195 to cfdf874 Compare April 14, 2026 07:58
@jopemachine jopemachine added the skip:changelog Make the action workflow to skip towncrier check label Apr 14, 2026
@jopemachine jopemachine removed the skip:changelog Make the action workflow to skip towncrier check label Apr 14, 2026
@jopemachine
Copy link
Copy Markdown
Member Author

Consolidated into the merged PR stack (A+B+C+D, E+F, G+H, I)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:docs Documentations comp:manager Related to Manager component size:M 30~100 LoC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants